#include "txr_algo_dlm_seg_demo.h"
#include <QFileDialog>
#include <QTextStream>
#include <QDir>
#include <QFile>
#include <sstream>
#include <iomanip> 
#include "../txr_algo_dlm_seg/txr_algo_dlm_seg.h"
#include "../tk_utility/tk_timer.hpp"
#include <QThread>
using namespace txr_algo_dlm_seg;

#include <opencv2\highgui.hpp>
#include <opencv2\imgproc.hpp>

const static int mark_color[txr_algo_dlm_seg_demo::DLM_CLASS_NUM][3] =
{
	{ 255, 0, 0 },
	{ 0, 255, 0 },
	{ 0, 0, 255 },

	{ 255, 0, 255 },
	{ 255, 255, 0 },
	{ 0, 255, 255 },

	{ 254, 67, 101 },
	{ 131, 175, 155 },
	{ 249, 205, 173 },

	{ 220, 87, 18 },
	{ 254, 67, 101 },
	{ 131, 175, 155 },

	{ 249, 205, 173 },
	{ 220, 87, 18 },
};
bool ReadImage(QString file_path, QVector<float> &v_data, int & row, int &col)
{
	using namespace cv;
	Mat src = imread(file_path.toLocal8Bit().data(), IMREAD_ANYCOLOR | IMREAD_ANYDEPTH);
	if (src.empty())
	{
		return false;
	}
	if (src.channels() != 3 && src.channels() != 1)
	{
		return false;
	}
	row = src.rows;
	col = src.cols;

	v_data.resize(row * col * 3);

	int img_size = row*col;
	if (src.channels() == 3)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3w color_value = src.ptr<Vec3w>(i)[j];
					v_data[i*col + j] = float(color_value[2]) / 65535.f;
					v_data[img_size + i*col + j] = float(color_value[1]) / 65535.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3b color_value = src.ptr<Vec3b>(i)[j];
					v_data[i*col + j] = float(color_value[2]) / 255.f;
					v_data[img_size + i*col + j] = float(color_value[1]) / 255.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}
	if (src.channels() == 1)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<ushort, 1> color_value = src.ptr<Vec<ushort, 1>>(i)[j];
					v_data[i*col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size + i*col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<uchar, 1> color_value = src.ptr<Vec<uchar, 1>>(i)[j];
					v_data[i*col + j] = float(color_value[0]) / 255.f;
					v_data[img_size + i*col + j] = float(color_value[0]) / 255.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}

	return true;
}


void MarkImage(st_detect_unit unit, QString dest_path, QVector<st_box> &result)
{
	bool * p_mask = unit.result.p_seg_mask;

	st_img_rgb img = unit.img;

	using namespace cv;
	Mat mat;

	mat.create(img.h, img.w, CV_8UC3);
	int img_size = img.h*img.w;
	for (int i = 0; i < img.h; i++)
	{
		for (int j = 0; j < img.w; j++)
		{
			mat.at<Vec3b>(i, j)[2] = img.pp_rgbf[0][i * img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[1] = img.pp_rgbf[1][i * img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[0] = img.pp_rgbf[2][i * img.w + j] * 255.f;
			if (p_mask && p_mask[i * img.w + j])
			{
				mat.at<Vec3b>(i, j)[2] = 255;
			}
		}
	}

	for (int i = 0; i < result.size(); ++i)
	{
		st_box r = result[i];

		cv::Rect rect(r.x- r.w/2, r.y- r.h/2, r.w, r.h);
		rectangle(mat, rect,
			Scalar(
				mark_color[r.obj_id][2],
				mark_color[r.obj_id][1],
				mark_color[r.obj_id][0]),
			1, 1, 0);

		st_polygon polygon = r.polygon;
		int node_num = polygon.node_num;
		for (auto j = 0; j < node_num; j++)
		{
			st_pt pt1 = polygon.pts[j];
			st_pt pt2 = polygon.pts[(j + 1) % node_num];
			cv::Point _pt1(pt1.x, pt1.y);
			cv::Point _pt2(pt2.x, pt2.y);
			cv::line(mat, _pt1, _pt2, cv::Scalar(0, 255, 0), 2);
		}

	}



	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}

	imwrite(dest_path.toLocal8Bit().data(), mat);
}

void SaveTxt(QString dest_path, QVector<st_box> &result)
{
	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}

	QFile::remove(dest_path);

	QFile file(dest_path);
	if (file.open(QIODevice::WriteOnly))
	{
		QTextStream stream(&file);
		for (int i = 0; i < result.size(); ++i)
		{
			stream << result[i].obj_id << " ";
			stream << result[i].prob << " ";
			stream << result[i].x << " ";
			stream << result[i].y << " ";
			stream << result[i].w << " ";
			stream << result[i].h << "\n";
		}
	}
	file.close();
}


txr_algo_dlm_seg_demo::txr_algo_dlm_seg_demo(QWidget *parent)
    : QMainWindow(parent)
{
    ui.setupUi(this);
	connect(ui.pushButtonOpenTestDir, SIGNAL(clicked()), this, SLOT(OnOpenTest()));
	connect(ui.pushButtonOpenModelDir, SIGNAL(clicked()), this, SLOT(OnOpenModel()));
	connect(ui.pushButtonTest, SIGNAL(clicked()), this, SLOT(OnTest()));

	checkBoxClass[0] = ui.checkBoxClass0; spinBoxTh[0] = ui.spinBoxTh0;
	checkBoxClass[1] = ui.checkBoxClass1; spinBoxTh[1] = ui.spinBoxTh1;
	checkBoxClass[2] = ui.checkBoxClass2; spinBoxTh[2] = ui.spinBoxTh2;
	checkBoxClass[3] = ui.checkBoxClass3; spinBoxTh[3] = ui.spinBoxTh3;
	checkBoxClass[4] = ui.checkBoxClass4; spinBoxTh[4] = ui.spinBoxTh4;
	checkBoxClass[5] = ui.checkBoxClass5; spinBoxTh[5] = ui.spinBoxTh5;
	checkBoxClass[6] = ui.checkBoxClass6; spinBoxTh[6] = ui.spinBoxTh6;
	checkBoxClass[7] = ui.checkBoxClass7; spinBoxTh[7] = ui.spinBoxTh7;
	checkBoxClass[8] = ui.checkBoxClass8; spinBoxTh[8] = ui.spinBoxTh8;
	checkBoxClass[9] = ui.checkBoxClass9; spinBoxTh[9] = ui.spinBoxTh9;
	checkBoxClass[10] = ui.checkBoxClass10; spinBoxTh[10] = ui.spinBoxTh10;
	checkBoxClass[11] = ui.checkBoxClass11; spinBoxTh[11] = ui.spinBoxTh11;
	checkBoxClass[12] = ui.checkBoxClass12; spinBoxTh[12] = ui.spinBoxTh12;
	checkBoxClass[13] = ui.checkBoxClass13; spinBoxTh[13] = ui.spinBoxTh13;
	checkBoxClass[14] = ui.checkBoxClass14; spinBoxTh[14] = ui.spinBoxTh14;
	for (tk_int32 i = 15; i < DLM_CLASS_NUM; ++i)
	{
		checkBoxClass[i] = ui.checkBoxClass14; spinBoxTh[i] = ui.spinBoxTh14;
	}

	ui.progressBar->setRange(0, 100);
	ui.progressBar->setValue(0);
}

void txr_algo_dlm_seg_demo::OnOpenTest()
{
	QString dir = QFileDialog::getExistingDirectory(this);
	if (dir.isEmpty())
	{
		return;
	}
	ui.labelDirTest->setText(dir);
}

void txr_algo_dlm_seg_demo::OnOpenModel()
{
	QString dir_path = QFileDialog::getExistingDirectory(this);
	if (dir_path.isEmpty())
	{
		return;
	}
	ui.labelDirModel->setText(dir_path);

	QDir dir(dir_path);
	QStringList filters;
	// 支持三种模型
	filters << "*.seg1" << "*.seg2" << "*.oseg2";

	QStringList files = dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	ui.comboBoxModel->clear();
	ui.comboBoxModel->addItems(files);
}

void txr_algo_dlm_seg_demo::OnTest()
{
	st_test_info info;

	info.test_dir = ui.labelDirTest->text();

	info.model_path = ui.labelDirModel->text() + "\\" + ui.comboBoxModel->currentText();

	info.save_ng = ui.checkBoxSaveNg->isChecked();
	info.save_gd = ui.checkBoxSaveGd->isChecked();

	for (int i = 0; i < DLM_CLASS_NUM; ++i)
	{
		info.class_en[i] = checkBoxClass[i]->isChecked();
		info.class_th[i] = spinBoxTh[i]->value();
	}

	DlmTest(info);
}

void txr_algo_dlm_seg_demo::DlmTest(st_test_info &info)
{
	using namespace txr_algo_dlm_seg;
	TxrAlgoDlmSeg dlm;
	{
		st_dlm_cfg cfg;

		QFileInfo model_file = QFileInfo(info.model_path);
		QString model_name = model_file.absoluteDir().path() + "\\" + model_file.baseName();
		QString model_suffix = model_file.suffix();

		sprintf_s(cfg.file_path_name, MAX_PATH_BYTES, "%s", model_name.toLocal8Bit().data());
		if (model_suffix == "seg1")
		{
			cfg.type = DLM_SEG_TYPE_YOLOV5;
		}
		else if (model_suffix == "seg2")
		{
			cfg.type = DLM_SEG_TYPE_DB;
		}
		else if (model_suffix == "oseg2")
		{
			cfg.type = DLM_SEG_TYPE_DB_O;
		}
		if (!dlm.Load(cfg))
		{
			return;
		}
	}
	//============================================================
	{
		QDir dir(info.test_dir + "\\result");
		dir.removeRecursively();
	}
	//============================================================

	QStringList test_files;
	{
		QDir dir(info.test_dir);
		QStringList filters;
		filters << "*.png" << "*.bmp";
		test_files =
			dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	}

	int cls_ng_box_count[DLM_CLASS_NUM] = { 0 };
	int img_gd_count = 0;
	int img_ng_count = 0;
	int file_count = 0;
	ui.progressBar->setValue(0);
	for each (QString file in test_files)
	{
		QString img_path = info.test_dir + "\\" + file;
		QVector<float> v_data;
		int row, col;
		if (!ReadImage(img_path, v_data, row, col))
		{
			continue;
		}

		QVector<bool> v_mask;
		v_mask.fill(false, row*col);

		txr_algo_dlm_seg::st_detect_unit unit;
		unit.img.w = col;
		unit.img.h = row;
		unit.img.c = 3;
		unit.img.pp_rgbf[0] = v_data.data() + col*row * 0;
		unit.img.pp_rgbf[1] = v_data.data() + col*row * 1;
		unit.img.pp_rgbf[2] = v_data.data() + col*row * 2;
		unit.param.prob_thres = 0.05;

		unit.result.p_seg_mask = v_mask.data();
		{
			tk::timer::universal_short_timer timer;
			dlm.Detect(&unit, 1);
			printf("dlm detect use time: %f ms \n", timer.elapsed_ms());
		}

		st_box * p_box = unit.result.boxes;
		tk_int32 box_size = unit.result.vaild_num;
		bool is_ng_flag = false;
		QVector<st_box> v_result_box;
		for (int i = 0; i < box_size; ++i)
		{
			int cls = p_box[i].obj_id;
			float prob = p_box[i].prob;
			if (info.class_en[cls] && prob >= float(info.class_th[cls]) / 100.f)
			{
				is_ng_flag = true;
				cls_ng_box_count[cls]++;
				v_result_box.push_back(p_box[i]);
			}
		}
		is_ng_flag ? img_ng_count += 1 : img_gd_count += 1;

		QFileInfo file_info = QFileInfo(img_path);
		if (info.save_gd && !is_ng_flag)
		{
			QString mark_file = info.test_dir + QString("\\result\\gd\\") + file_info.baseName() + ".bmp";
			MarkImage(unit, mark_file, v_result_box);
		}
		if (info.save_ng && is_ng_flag)
		{
			QString mark_file = info.test_dir + QString("\\result\\ng\\") + file_info.baseName() + ".bmp";
			MarkImage(unit, mark_file, v_result_box);
		}
		if (info.save_txt && is_ng_flag)
		{
			QString mark_file = info.test_dir + QString("\\result\\txt\\") + file_info.baseName() + ".txt";
			SaveTxt(mark_file, v_result_box);
		}
		file_count++;
		ui.progressBar->setValue(file_count * 100 / test_files.size());

		printf("total count : %d good count: %d  ng count: %d \n", img_gd_count + img_ng_count, img_gd_count, img_ng_count);
		printf("class: ");
		for (int i = 0; i < DLM_CLASS_NUM; ++i)
		{
			printf("[%d] = %d ,", i, cls_ng_box_count[i]);
		}
		printf("\n");

		QCoreApplication::processEvents();
	}
	ui.progressBar->setValue(100);
	printf("done!==========================================================\n");
}